package com.hc.app.manager;

import com.hc.app.dto.LeakBucketDTO;
import org.springframework.stereotype.Component;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.locks.ReentrantLock;

/**
 * @Author: huangcheng
 * @CreateTime: 2024-07-25  14:58
 * @Description:
 * @Version: 1.0
 */
@Component
public class LeakBucketManager {
    private final Map<String, LeakBucketDTO> leakBucketMap = new ConcurrentHashMap<>();
    private final ReentrantLock lock = new ReentrantLock();

    public void createLeakBucket(String key, long capacity, double refillRate) {
        lock.lock();
        try {
            if (!leakBucketMap.containsKey(key)) {
                LeakBucketDTO leakBucket = new LeakBucketDTO(capacity, refillRate);
                leakBucketMap.put(key, leakBucket);
            }
        } finally {
            lock.unlock();
        }
    }

    /**
     * 是否可以获取令牌
     * @param key
     * @param permits
     * @return
     */
    public boolean tryAcquire(String key, int permits) {
        LeakBucketDTO tokenBucket = leakBucketMap.get(key);
        if (tokenBucket != null) {
            return tokenBucket.tryAcquire(permits);
        }
        return false;
    }

    public void deductionToken(String key, int permits) {
        LeakBucketDTO leakBucketDTO = leakBucketMap.get(key);
        if (leakBucketDTO != null) {
            leakBucketDTO.deductionToken(permits);
        }
    }
}
